#!/usr/bin/env python3
# coding=utf-8

import rospy
from sensor_msgs.msg import JointState
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# Global data buffers
time_vals = []
positions_vals = [[] for _ in range(7)]  # Actual joint positions
cmd_positions_vals = [[] for _ in range(7)]  # Commanded joint positions
has_cmd_data = False  # Flag to track if we've received any command data

def joint_callback(msg):
    """ROS Callback for actual joint states."""
    if not time_vals:
        time_vals.append(0.0)
    else:
        time_vals.append(time_vals[-1] + 0.05)  # Assuming 20Hz data
    
    for i in range(7):
        if i < len(msg.position):
            positions_vals[i].append(msg.position[i])
        else:
            positions_vals[i].append(0.0)  # Default value if position not available

def cmd_callback(msg):
    """ROS Callback for commanded joint states."""
    global has_cmd_data
    has_cmd_data = True
    
    for i in range(7):
        if i < len(msg.position):
            cmd_positions_vals[i].append(msg.position[i])
        else:
            cmd_positions_vals[i].append(0.0)  # Default value if position not available

def animate(i):
    """Function called periodically by FuncAnimation to update plots."""
    if not time_vals:
        return
    
    for idx, ax in enumerate(axs.flatten()):
        ax.cla()
        
        if idx < 7:
            # Always plot actual joint data
            ax.plot(time_vals, positions_vals[idx], 'b-', label=f'Actual Joint {idx+1}')
            
            # Only plot commanded data if we have it
            if has_cmd_data and len(cmd_positions_vals[idx]) == len(time_vals):
                ax.plot(time_vals, cmd_positions_vals[idx], 'r--', label=f'Commanded Joint {idx+1}')
            
            ax.set_xlabel('Time (sec)')
            ax.set_ylabel('Position (radian)')
            ax.legend()
            ax.grid()
        else:
            ax.axis('off')  # Hide the extra subplot

def main():
    """Main function."""
    rospy.init_node('joints_plotter', anonymous=True)
    
    # Subscribe to both topics
    rospy.Subscriber('/joint_states', JointState, joint_callback)
    rospy.Subscriber('/cmd_states', JointState, cmd_callback)

    # Set up matplotlib
    global fig, axs
    fig, axs = plt.subplots(4, 2, figsize=(12, 14))
    fig.suptitle('Joint Positions: Actual vs Commanded', fontsize=16)
    
    # Adjust layout to prevent overlap
    plt.tight_layout(rect=[0, 0, 1, 0.96])  # Leave space for suptitle
    
    # Start animation (50ms refresh rate)
    ani = animation.FuncAnimation(fig, animate, interval=50)
    
    plt.show()
    rospy.spin()

if __name__ == '__main__':
    main()